#include "ecc.h"

const curve sm2 = {
        .p = {0xffffffff, 0xffffffff, 0x00000000, 0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff, 0xfffffffe},
        .a = {0xfffffffc, 0xffffffff, 0x00000000, 0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff, 0xfffffffe},
        .b = {0x4d940e93, 0xddbcbd41, 0x15ab8f92, 0xf39789f5, 0xcf6509a7, 0x4d5a9e4b, 0x9d9f5e34, 0x28e9fa9e},
        .n = {0x39d54123, 0x53bbf409, 0x21c6052b, 0x7203df6b, 0xffffffff, 0xffffffff, 0xffffffff, 0xfffffffe},
        .g.x = {0x334c74c7, 0x715a4589, 0xf2660be1, 0x8fe30bbf, 0x6a39c994, 0x5f990446, 0x1f198119, 0x32c4ae2c},
        .g.y = {0x2139f0a0, 0x02df32e5, 0xc62a4740, 0xd0a9877c, 0x6b692153, 0x59bdcee3, 0xf4f6779c, 0xbc3736a2}
};

void point_fm_fp(point *p, const fp x, const fp y) {
    fp_set(p->x, x);
    fp_set(p->y, y);
}

void point_fm_str(point *p, const str x, const str y) {
    fp_fm_str(p->x, x);
    fp_fm_str(p->y, y);
}

void point_fm_byt(point *p, const byt x, const byt y) {
    fp_fm_byt(p->x, x);
    fp_fm_byt(p->y, y);
}

void point_fm_compress_octets(point *p, const u8 *src) {
    fp t1, t2;
    fp_fm_byt(p->x, src + 1);
    fp_tri(t1, p->x);
    fp_mul(t2, p->x, sm2.a);
    fp_add(t1, t1, t2);
    fp_add(t1, t1, sm2.b);
    // fp_sqrt(p->y, t1, e.p);
    if (fp_is_even(p->y) && src[0] == 0x03)
        fp_neg(p->y, p->y);
    if (fp_is_odd(p->y) && src[0] == 0x02)
        fp_neg(p->y, p->y);
    // if (point_is_on_curve(p) != 1)
    //     error_print;
}

void point_fm_octets(point *p, const u8 *in) {
    if (in[0] >= 0x04 && in[0] <= 0x07) {
        fp_fm_byt(p->x, in + 1);
        fp_fm_byt(p->y, in + 33);
    } else if (in[0] == 0x02 || in[0] == 0x03) {
        point_fm_compress_octets(p, in);
    } else {
        return;
    }
}

void point_fm_der(point *p, const u8 **src, u64 *slen) {
    u8 octets[65];
    u64 len;
    asn1_bit_string_from_der(octets, &len, src, slen);
    if (len != 65 || octets[0] != 0x04) {
        point_fm_fp(p, zero, zero);
    } else
        point_fm_octets(p, octets);
}

void point_to_compress_octets(point *p, u8 dst[33]) {
    *dst++ = fp_is_even(p->y) ? 0x02 : 0x03;
    fp_to_byt(p->x, dst);
}

void point_to_uncompressed_octets(point *p, u8 dst[65]) {
    *dst++ = 0x04;
    fp_to_byt(p->x, dst);
    fp_to_byt(p->y, dst + 32);
}

void point_to_mix_octets(point *p, u8 dst[65]) {
    *dst++ = fp_is_even(p->y) ? 0x06 : 0x07;
    fp_to_byt(p->x, dst);
    fp_to_byt(p->y, dst + 32);
}

int point_is_on_curve(point *r) {
    fp t1, t2, t3;
    fp_sqr(t1, r->y);
    fp_tri(t2, r->x);
    fp_mul(t3, r->x, sm2.a);
    fp_add(t2, t2, t3);
    fp_add(t2, t2, sm2.b);
    return fp_cmp(t1, t2);
}

void point_mul(point *p, fp k, const point *q) {
    jpoint g, t;
    jpoint_set(&g, q->x, q->y, one);
    jpoint_mul(&t, k, &g);
    point_fm_jpoint(p, &t);
}

void point_mul_fixed(point *r, fp k) {
    jpoint g;
    jpoint_mul_fixed(&g, k);
    point_fm_jpoint(r, &g);
}

void point_mutmul(point *r, fp s, fp t, const point *q) {
    jpoint g, t1, t2;
    jpoint_mul_fixed(&t1, s);
    jpoint_set(&g, q->x, q->y, one);
    jpoint_mul(&t2, t, &g);
    jpoint_add(&g, &t1, &t2);
    point_fm_jpoint(r, &g);
}

void point_to_jpoint(jpoint *r, const point *p) {
    jpoint_set(r, p->x, p->y, one);
}

void point_fm_jpoint(point *r, const jpoint *a) {
    fp tmp = {0}, z_inv = {0};
    fp_inv(z_inv, a->z);
    fp_mul(tmp, z_inv, z_inv);
    fp_mul(r->x, tmp, a->x);
    fp_mul(tmp, tmp, z_inv);
    fp_mul(r->y, tmp, a->y);
}

void inline jpoint_set(jpoint *p, const fp x, const fp y, const fp z) {
    fp_set(p->x, x);
    fp_set(p->y, y);
    fp_set(p->z, z);
}

void jpoint_add(jpoint *r, const jpoint *p, const jpoint *q) {
    if (jpoint_is_at_infinity(p)) {
        jpoint_set(r, q->x, q->y, q->z); return;
    }
    if (jpoint_is_at_infinity(q)) {
        jpoint_set(r, p->x, p->y, p->z); return;
    }
    const u64 *x1 = p->x, *y1 = p->y, *z1 = p->z;
    const u64 *x2 = q->x, *y2 = q->y, *z2 = q->z;
    fp x3, y3, z3, u1, u2, s1, s2, h, n, h2, h3, u1h2, t1, t2;
    fp_mul(t1, z1, z1);                 // t1 = z1^2
    fp_mul(t2, z2, z2);                 // t2 = z2^2
    fp_mul(u1, x1, t2);                 // u1 = x1 * z2^2
    fp_mul(u2, x2, t1);                 // u2 = x2 * z1^2
    fp_mul(t1, t1, z1);                 // t1 = z1^3
    fp_mul(t2, t2, z2);                 // t2 = z2^3
    fp_mul(s1, y1, t2);                 // s1 = y1 * z2^3
    fp_mul(s2, y2, t1);                 // s2 = y2 * z1^3
    if (fp_equ(u1, u2)) {
        if (fp_equ(s1, s2))
            jpoint_dou(r, p);
        else
            jpoint_set(r, zero, zero, one);
        return;
    }
    fp_sub(h, u2, u1);                  // h = u2 - u1
    fp_sub(n, s2, s1);                  // n = s2 - s1
    fp_mul(h2, h, h);                   // h2 = h^2
    fp_mul(h3, h2, h);                  // h3 = h^3
    fp_mul(u1h2, u1, h2);               // u1h2 = u1 * h^2
    fp_dou(t1, u1h2);                   // t1 = 2u1h2
    fp_mul(x3, n, n);                   // x3 = n^2
    fp_sub(x3, x3, h3);                 // x3 = n^2 - h3
    fp_sub(x3, x3, t1);                 // x3 = n^2 - h3 - 2u1h2
    fp_mul(t1, s1, h3);                 // t1 = s1 * h3
    fp_sub(y3, u1h2, x3);               // y3 = u1h2 - x3
    fp_mul(y3, y3, n);                  // y3 = n * (u1h2 - x3)
    fp_sub(y3, y3, t1);                 // y3 = n * (u1h2 - x3) - s1 * h3
    fp_mul(z3, z1, z2);                 // z3 = z1 * z2
    fp_mul(z3, z3, h);                  // z3 = h * z1 * z2
    jpoint_set(r, x3, y3, z3);
}

void jpoint_sub(jpoint *r, const jpoint *p, const jpoint *q) {
    fp y2, x3, y3, z3, u1, u2, s1, s2, h, n, h2, h3, u1h2, t1, t2;
    const u64 *x1 = p->x, *y1 = p->y, *z1 = p->z;
    const u64 *x2 = q->x, *z2 = q->z;
    fp_neg(y2, q->y);
    if (jpoint_is_at_infinity(p)) {
        jpoint_set(r, x2, y2, z2); return;
    }
    if (jpoint_is_at_infinity(q)) {
        jpoint_set(r, x1, y1, z1); return;
    }

    fp_mul(t1, z1, z1);                 // t1 = z1^2
    fp_mul(t2, z2, z2);                 // t2 = z2^2
    fp_mul(u1, x1, t2);                 // u1 = x1 * z2^2
    fp_mul(u2, x2, t1);                 // u2 = x2 * z1^2
    fp_mul(t1, t1, z1);                 // t1 = z1^3
    fp_mul(t2, t2, z2);                 // t2 = z2^3
    fp_mul(s1, y1, t2);                 // s1 = y1 * z2^3
    fp_mul(s2, y2, t1);                 // s2 = y2 * z1^3
    if (fp_equ(u1, u2)) {
        if (fp_equ(s1, s2))
            jpoint_dou(r, p);
        else
            jpoint_set(r, zero, zero, one);
        return;
    }
    fp_sub(h, u2, u1);                  // h = u2 - u1
    fp_sub(n, s2, s1);                  // n = s2 - s1
    fp_mul(h2, h, h);                   // h2 = h^2
    fp_mul(h3, h2, h);                  // h3 = h^3
    fp_mul(u1h2, u1, h2);               // u1h2 = u1 * h^2
    fp_dou(t1, u1h2);                   // t1 = 2u1h2
    fp_mul(x3, n, n);                   // x3 = n^2
    fp_sub(x3, x3, h3);                 // x3 = n^2 - h3
    fp_sub(x3, x3, t1);                 // x3 = n^2 - h3 - 2u1h2
    fp_mul(t1, s1, h3);                 // t1 = s1 * h3
    fp_sub(y3, u1h2, x3);               // y3 = u1h2 - x3
    fp_mul(y3, y3, n);                  // y3 = n * (u1h2 - x3)
    fp_sub(y3, y3, t1);                 // y3 = n * (u1h2 - x3) - s1 * h3
    fp_mul(z3, z1, z2);                 // z3 = z1 * z2
    fp_mul(z3, z3, h);                  // z3 = h * z1 * z2
    jpoint_set(r, x3, y3, z3);
}

void jpoint_dou(jpoint *r, const jpoint *a) {
    if (jpoint_is_at_infinity(a)) {
        jpoint_copy(r, a); return;
    }
    // A = 3(x1 - z1^2) * (x1 + z1^2)
    // B = 2Y1, z3 = B * z1, C = B^2
    // D = C * x1, x3 = A^2 - 2D, (D - x3) * A - C^2/2
    const u64 *x1 = a->x, *y1 = a->y, *z1 = a->z;
    fp t1, t2, t3, x3, y3, z3;
    fp_mul(t1, z1, z1);                 // t1 = z1^2
    fp_sub(t2, x1, t1);                 // t2 = x1 - z1^2
    fp_add(t1, x1, t1);                 // t1 = x1 + z1^2
    fp_mul(t2, t2, t1);                 // t2 = x1^2 - z1^4
    fp_dou(t3, t2);                     // t3 = 2(x1^2 - z1^4)
    fp_add(t2, t2, t3);                 // t2 = A = 3t2 = 3(x1^2 - z1^4)
    fp_dou(y3, y1);                     // y3 = B = 2y1
    fp_mul(z3, y3, z1);                 // z3 = B * z1
    fp_mul(y3, y3, y3);                 // y3 = C = B^2
    fp_mul(t3, y3, x1);                 // t3 = D = C * x1
    fp_mul(y3, y3, y3);                 // y3 = C^2
    fp_haf(y3, y3);                     // y3 = C^2/2
    fp_mul(x3, t2, t2);                 // x3 = A^2
    fp_dou(t1, t3);                     // t1 = 2D
    fp_sub(x3, x3, t1);                 // x3 = A^2 - 2D
    fp_sub(t1, t3, x3);                 // t1 = D - x3
    fp_mul(t1, t1, t2);                 // t1 = (D - x3) * A
    fp_sub(y3, t1, y3);                 // y3 = (D - x3) * A - C^2/2
    jpoint_set(r, x3, y3, z3);
}

void jpoint_mdo(jpoint *r, u32 m, const jpoint *p) {
    if (jpoint_is_at_infinity(p)) {
        jpoint_copy(r, p); return;
    }
    fp x, y, z, w, a, b, t;
    fp_set(x, p->x);
    fp_set(y, p->y);
    fp_set(z, p->z);
    fp_dou(y, y);                // y = 2Y
    fp_mul(w, z, z);
    fp_mul(w, w, w);             // w = Z^4
    while (m > 0) {
        fp_mul(a, x, x);
        fp_sub(a, a, w);
        fp_dou(t, a);
        fp_add(a, a, t);        // a = 3(x^2 -w)
        fp_mul(b, y, y);
        fp_mul(b, x, b);        // b = x * y^2
        fp_mul(x, a, a);
        fp_sub(x, x, b);
        fp_sub(x, x, b);        // x = a^2 - 2b
        fp_mul(z, z, y);        // z = z * y
        m--;
        fp_mul(t, y, y);
        fp_mul(t, t, t);           // t = y^4
        if (m != 0)
            fp_mul(w, w, t);    // w = w * y^4
        fp_sub(y, b, x);
        fp_mul(y, y, a);
        fp_dou(y, y);
        fp_sub(y, y, t);        // y = 2(b-x) - y^4
    }
    fp_haf(y, y);
    jpoint_set(r, x, y, z);
}

void jpoint_mul(jpoint *r, fp k, const jpoint *g) {
    naf kn;
    u32 i = 0, j = 1;
    jpoint upt[8];
    fp_naf(kn, 4, k);
    jpoint_set(&upt[0], g->x, g->y, g->z);
    jpoint_dou(r, g);
    jpoint_add(&upt[1], &upt[0], r);
    jpoint_add(&upt[2], &upt[1], r);
    jpoint_add(&upt[3], &upt[2], r);
    jpoint_add(&upt[4], &upt[3], r);
    jpoint_add(&upt[5], &upt[4], r);
    jpoint_add(&upt[6], &upt[5], r);
    jpoint_add(&upt[7], &upt[6], r);
    jpoint_set_infinity(r);
    do {
        if (kn[i] == 0) {
            j++;
        } else {
            jpoint_mdo(r, j, r);
            if (kn[i] > 0)
                jpoint_add(r, r, &upt[kn[i] >> 1]);
            else
                jpoint_sub(r, r, &upt[(-kn[i]) >> 1]);
            j = 1;
        }
    } while (++i <= 256);
    jpoint_mdo(r, j - 1, r);
    Clean(kn);
}

void jpoint_mul_fixed(jpoint *r, fp k) {
    extern const point fpt[52];
    jpoint a, b;
    jpoint_set_infinity(&a);
    jpoint_set_infinity(&b);

    naf kn;
    fp_naf(kn, 2, k);
    int K[52];
    int i, j = 0;
    for (i = 256; i > 2; i = i-5){
        K[j++] = kn[i] + (kn[i-1]<<1) + (kn[i-2]<<2) + (kn[i-3]<<3) + (kn[i-4]<<4);
    }
    K[j] = kn[1] + (kn[0]<<1);

    for (j = 21; j > 0; j--) {
        for (i = 0; i < 52; i++) {
            if (K[i] == j)
                jpoint_add_AJ(&b, &b, &fpt[i]);
            if (K[i] == -j)
                jpoint_sub_AJ(&b, &b, &fpt[i]);
        }
        jpoint_add(&a, &a, &b);
    }
    jpoint_set(r, a.x, a.y, a.z);
    Clean(kn);
}

void jpoint_add_AJ(jpoint *r, const jpoint *p, const point *q) {
    if (jpoint_is_at_infinity(p)) {
        jpoint_set(r, q->x, q->y, one); return;
    }
    if (fp_equ(q->x, zero) && fp_equ(q->y, zero)) {
        jpoint_set(r, p->x, p->y, p->z); return;
    }

    const u64 *x1 = p->x, *y1 = p->y, *z1 = p->z, *x2 = q->x, *y2 = q->y;
    fp x3, y3, z3, t1, t2, t3, t4;
    fp_mul(t1, z1, z1);                 // t1 = A = z1^2
    fp_mul(t2, t1, z1);                 // t2 = B = z1 * A
    fp_mul(t1, t1, x2);                 // t1 = C = x2 * A
    fp_mul(t2, t2, y2);                 // t2 = D = y2 * B
    fp_sub(t1, t1, x1);                 // t1 = E = C - x1
    fp_sub(t2, t2, y1);                 // t2 = F = D - y1
    if (fp_is_0(t1)) {
        if (fp_is_0(t2)) {
            jpoint t;
            jpoint_set(&t, x2, y2, one);
            jpoint_dou(r, &t);
        } else
            jpoint_set_infinity(r);
        return;
    }
    fp_mul(z3, z1, t1);                 // z3 = z1 * E
    fp_mul(t3, t1, t1);                 // t3 = G = E^2
    fp_mul(t4, t3, t1);                 // t4 = H = E^3
    fp_mul(t3, t3, x1);                 // t3 = I = x1 * G
    fp_dou(t1, t3);                     // t1 = 2I
    fp_mul(x3, t2, t2);                 // x3 = F^2
    fp_sub(x3, x3, t1);                 // x3 = F^2 - 2I
    fp_sub(x3, x3, t4);                 // x3 = F^2 - 2I - H
    fp_sub(t3, t3, x3);                 // t3 = I - x3
    fp_mul(t3, t3, t2);                 // t3 = (I - x3) * F
    fp_mul(t4, t4, y1);                 // t4 = y1 * H
    fp_sub(y3, t3, t4);                 // y3 = (I - x3) * F - y1 * H
    jpoint_set(r, x3, y3, z3);
}

void jpoint_sub_AJ(jpoint *r, const jpoint *p, const point *q) {
    const u64 *x1 = p->x, *y1 = p->y, *z1 = p->z, *x2 = q->x;
    fp y2, x3, y3, z3, t1, t2, t3, t4;
    fp_neg(y2, q->y);

    if (jpoint_is_at_infinity(p)) {
        jpoint_set(r, x2, y2, one); return;
    }
    if (fp_equ(q->x, zero) && fp_equ(q->y, zero)) {
        jpoint_set(r, x1, y1, z1); return;
    }
    fp_mul(t1, z1, z1);                 // t1 = A = z1^2
    fp_mul(t2, t1, z1);                 // t2 = B = z1 * A
    fp_mul(t1, t1, x2);                 // t1 = C = x2 * A
    fp_mul(t2, t2, y2);                 // t2 = D = y2 * B
    fp_sub(t1, t1, x1);                 // t1 = E = C - x1
    fp_sub(t2, t2, y1);                 // t2 = F = D - y1
    if (fp_is_0(t1)) {
        if (fp_is_0(t2)) {
            jpoint t;
            jpoint_set(&t, x2, y2, one);
            jpoint_dou(r, &t);
        } else
            jpoint_set_infinity(r);
        return;
    }
    fp_mul(z3, z1, t1);                 // z3 = z1 * E
    fp_mul(t3, t1, t1);                 // t3 = G = E^2
    fp_mul(t4, t3, t1);                 // t4 = H = E^3
    fp_mul(t3, t3, x1);                 // t3 = I = x1 * G
    fp_dou(t1, t3);                     // t1 = 2I
    fp_mul(x3, t2, t2);                 // x3 = F^2
    fp_sub(x3, x3, t1);                 // x3 = F^2 - 2I
    fp_sub(x3, x3, t4);                 // x3 = F^2 - 2I - H
    fp_sub(t3, t3, x3);                 // t3 = I - x3
    fp_mul(t3, t3, t2);                 // t3 = (I - x3) * F
    fp_mul(t4, t4, y1);                 // t4 = y1 * H
    fp_sub(y3, t3, t4);                 // y3 = (I - x3) * F - y1 * H
    jpoint_set(r, x3, y3, z3);
}

